package com.vr.xxaiagent.chatmemory;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import org.objenesis.strategy.StdInstantiatorStrategy;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

/**
 * 基于文件的持久化记忆
 *
 * @author hzh
 * @date 2025/06/01
 */
public class FileBaseChatMemory implements ChatMemory {
    private final String BASE_DIR;
    private static final Kryo kryo = new Kryo();

    static {
        // 关闭手动注册，默认为true
        kryo.setRegistrationRequired(false);
        // 设置实例化策略,使用标准的实例化策略
        kryo.setInstantiatorStrategy(new StdInstantiatorStrategy());
    }

    public FileBaseChatMemory(String dir) {
        this.BASE_DIR = dir;
        File baseDir = new File(BASE_DIR);
        if (!baseDir.exists()) {
            // 文件夹不存在时创建
            baseDir.mkdirs();
        }
    }

    @Override
    public void add(String conversationId, List<Message> messages) {
        // 添加消息
        // 基于会话id先获取之前的会话消息list
        List<Message> conversationMessageList = getOrCreateConversation(conversationId);
        // 再把新的消息添加到会话消息list中
        conversationMessageList.addAll(messages);
        // 再重新保存(序列化)
        saveConversation(conversationId, conversationMessageList);
    }




    @Override
    public List<Message> get(String conversationId, int lastN) {
        // 拿到会话消息list
        List<Message> allMessages = getOrCreateConversation(conversationId);
        return allMessages.stream()
                // 10-3=7:假如共有10个元素，获取最新的是三个，取倒数第三个及以后，也就是跳到第七个
                .skip(Math.max(0, allMessages.size() - lastN))
                .collect(Collectors.toList());
    }

    @Override
    public void clear(String conversationId) {
        File file = getConversationFile(conversationId);
        if (file.exists()) {
            file.delete();
        }
    }

    private List<Message> getOrCreateConversation(String conversationId) {
        // 获取文件,存在时读取
        File file = getConversationFile(conversationId);
        List<Message> conversationMessageList = new ArrayList<>();
        if (file.exists()) {
            try (Input input = new Input(new FileInputStream(file))) {
                conversationMessageList = kryo.readObject(input, ArrayList.class);
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        return conversationMessageList;
    }

    private File getConversationFile(String conversationId) {
        // 存在时获取
        return new File(BASE_DIR, conversationId + ".kryo");
    }
    private void saveConversation(String conversationId, List<Message> conversationMessageList) {
        File conversationFile = getConversationFile(conversationId);
         try (Output output = new Output(new FileOutputStream(conversationFile))) {
            kryo.writeObject(output, conversationMessageList);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }
}
